package top.techqi.usbd.driver.serial

import android.hardware.usb.UsbConstants
import android.hardware.usb.UsbDevice
import android.hardware.usb.UsbInterface
import top.techqi.core.kext.WORD
import top.techqi.core.util.LogUtil
import java.io.IOException
import java.util.EnumSet

class Wch341SerialDriver(device: UsbDevice, intfId: Int) : BasicsSerialDriver(device, intfId) {
    private var version: Byte? = null
    private var baudRate: Int = DEFAULT_BAUD_RATE

    override val supportedControlLines: EnumSet<ControlLine> =
        EnumSet.allOf(ControlLine::class.java)

    override val cd: Boolean
        get() = (getStatus() and GET_BIT_CD) == 0

    override val ri: Boolean
        get() = (getStatus() and GET_BIT_RI) == 0

    override val cts: Boolean
        get() = (getStatus() and GET_BIT_CTS) == 0

    override val dsr: Boolean
        get() = (getStatus() and GET_BIT_DSR) == 0

    override var dtr: Boolean = false
        set(value) {
            field = value
            setModemCtrl()
        }

    override var rts: Boolean = false
        set(value) {
            field = value
            setModemCtrl()
        }

    @Throws(IOException::class)
    override fun openInterfaces() {
        for (index in 0 until device.interfaceCount) {
            if (!conn.claimInterface(device.getInterface(index), true)) {
                throw IOException("Could not claim data interface")
            }
        }

        val intf: UsbInterface = device.getInterface(device.interfaceCount - 1)
        for (i in 0 until intf.endpointCount) {
            val ep = intf.getEndpoint(i)
            if (ep.type == UsbConstants.USB_ENDPOINT_XFER_BULK) {
                if (ep.direction == UsbConstants.USB_DIR_IN) {
                    rEndpoint = ep
                } else {
                    wEndpoint = ep
                }
            }
        }

        initialize()
    }

    override fun closeInterfaces() {
        for (index in 0 until device.interfaceCount) {
            conn.releaseInterface(device.getInterface(index))
        }
    }

    override fun setParameters(baudRate: Int, dataBits: Byte, stopBits: Byte, parity: Byte) {
        setBaudRate(baudRate)
        setLineCtrl(dataBits, stopBits, parity)
    }

    private fun controlRx(request: Int, value: Int, index: Int, buffer: ByteArray): Int {
        return conn.controlTransfer(
            UsbConstants.USB_TYPE_VENDOR or UsbConstants.USB_DIR_IN,
            request, value, index, buffer, buffer.size, DEFAULT_TIMEOUT_RW
        )
    }

    private fun controlTx(request: Int, value: Int, index: Int): Int {
        return conn.controlTransfer(
            UsbConstants.USB_TYPE_VENDOR or UsbConstants.USB_DIR_OUT,
            request, value, index, null, 0, DEFAULT_TIMEOUT_RW
        )
    }

    @Throws(IOException::class)
    private fun initialize() {
        val buffer = ByteArray(2)
        if (controlRx(REQ_READ_VERSION, 0, 0, buffer) < 0) {
            throw IOException("Error read chip version")
        } else {
            version = buffer[0]
            LogUtil.d(TAG, "Chip version: 0x%1$02X".format(version))
        }

        if (controlTx(REQ_SERIAL_INIT, 0, 0) < 0 ||
            controlRx(REQ_READ_REG, WORD(REG_LCR_DUE, REG_LCR_UNO), 0, buffer) < 0 ||
            controlTx(
                REQ_WRITE_REG, WORD(REG_LCR_DUE, REG_LCR_UNO),
                LCR_ENABLE_RX or LCR_ENABLE_TX or LCR_DATA_CS8
            ) < 0 ||
            controlTx(REQ_SERIAL_INIT, 0x501f, 0xd90a) < 0
        ) {
            throw IOException("Init serial failed")
        }

        setBaudRate(DEFAULT_BAUD_RATE)
        setModemCtrl()
    }

    @Throws(IOException::class)
    private fun setBaudRate(baudRate: Int) {
        if (baudRate <= 0) throw IllegalArgumentException("Invalid baud rate: $baudRate")

        var factor: Int
        var divisor: Int
        if (baudRate == 921600) {
            divisor = 7
            factor = 0xf300
        } else {
            factor = BAUDBASE_FACTOR / baudRate
            divisor = BAUDBASE_DIVMAX
            while ((factor > 0xfff0) && divisor > 0) {
                factor = factor shr 3
                divisor--
            }
            if (factor > 0xfff0) {
                throw UnsupportedOperationException("Unsupported baud rate: $baudRate")
            }
            factor = 0x10000 - factor
        }

        divisor = divisor or 0x0080 // else ch341a waits until buffer full
        val value0 = (factor and 0xff00) or divisor
        val value1 = (factor and 0x00ff)

        if (controlTx(REQ_WRITE_REG, WORD(REG_BPS_DIV, REG_BPS_PRE), value0) < 0 ||
            controlTx(REQ_WRITE_REG, 0x0f2c, value1) < 0
        ) {
            throw IOException("Error setting baud rate: $baudRate")
        } else {
            this.baudRate = baudRate
        }
    }

    @Throws(IOException::class)
    private fun setLineCtrl(dataBits: Byte, stopBits: Byte, parity: Byte) {
        val lineCtrlReg: Int = LCR_ENABLE_RX or LCR_ENABLE_TX or when (dataBits) {
            DATABITS_5 -> LCR_DATA_CS5
            DATABITS_6 -> LCR_DATA_CS6
            DATABITS_7 -> LCR_DATA_CS7
            DATABITS_8 -> LCR_DATA_CS8
            else -> throw IllegalArgumentException("Invalid data bits: $dataBits")
        } or when (stopBits) {
            STOPBITS_1 -> 0
            STOPBITS_2 -> LCR_STOP_BITS_2
            else -> throw IllegalArgumentException("Invalid stop bits: $stopBits")
        } or when (parity) {
            PARITY_ODD -> LCR_ENABLE_PAR
            PARITY_EVEN -> LCR_ENABLE_PAR or LCR_PAR_EVEN
            PARITY_MARK -> LCR_ENABLE_PAR or LCR_MARK_SPACE
            PARITY_SPACE -> LCR_ENABLE_PAR or LCR_MARK_SPACE or LCR_PAR_EVEN
            PARITY_NONE -> 0
            else -> throw IllegalArgumentException("Invalid parity: $parity")
        }

        if (controlTx(REQ_WRITE_REG, WORD(REG_LCR_DUE, REG_LCR_UNO), lineCtrlReg) < 0) {
            throw IOException("Error setting ctrl byte")
        }
    }

    @Throws(IOException::class)
    private fun setModemCtrl() {
        val modemCtrl = (if (dtr) SET_BIT_DTR else 0) or (if (rts) SET_BIT_RTS else 0)
        if (controlTx(REQ_MODEM_CTRL, modemCtrl.inv(), 0) < 0) {
            throw IOException("Error setting DTR/RTS byte")
        }
    }

    @Throws(IOException::class)
    private fun getStatus(): Int {
        val buffer = ByteArray(2)
        val ret = controlRx(REQ_READ_REG, 0x0706, 0, buffer)
        if (ret < 0) throw IOException("Error getting ctrl byte")
        return buffer[0].toInt()
    }

    companion object {
        private const val TAG = "Ch341Driver"

        private const val DEFAULT_BAUD_RATE = 9600
        private const val DEFAULT_TIMEOUT_RW = 5000

        private const val BAUDBASE_FACTOR = 1532620800
        private const val BAUDBASE_DIVMAX = 3

        private const val REQ_READ_VERSION = 0x5f
        private const val REQ_WRITE_REG = 0x9a
        private const val REQ_READ_REG = 0x95
        private const val REQ_SERIAL_INIT = 0xa1
        private const val REQ_MODEM_CTRL = 0xa4

        private const val REG_BPS_PRE = 0x12
        private const val REG_BPS_DIV = 0x13
        private const val REG_LCR_UNO = 0x18
        private const val REG_LCR_DUE = 0x25

        private const val LCR_ENABLE_RX = 0x80
        private const val LCR_ENABLE_TX = 0x40
        private const val LCR_MARK_SPACE = 0x20
        private const val LCR_PAR_EVEN = 0x10
        private const val LCR_ENABLE_PAR = 0x08
        private const val LCR_STOP_BITS_2 = 0x04
        private const val LCR_DATA_CS8 = 0x03
        private const val LCR_DATA_CS7 = 0x02
        private const val LCR_DATA_CS6 = 0x01
        private const val LCR_DATA_CS5 = 0x00

        private const val GET_BIT_CTS = 0x01
        private const val GET_BIT_DSR = 0x02
        private const val GET_BIT_RI = 0x04
        private const val GET_BIT_CD = 0x08

        private const val SET_BIT_DTR = 0x20
        private const val SET_BIT_RTS = 0x40
    }
}
